import torch
import torch.nn as nn
from torch.nn import Module
from networks import MixedFeatureNet
import os
from timm.models.vision_transformer import VisionTransformer

class ViTEnhancedNet(nn.Module):
    def __init__(self, num_class=7, num_head=2, pretrained=True):
        super(ViTEnhancedNet, self).__init__()
        
        # 输入通道调整
        # self.input_adjust = nn.Conv2d(2, 3, kernel_size=1)
        
        # 加载预训练的ViT模型
        self.vit = VisionTransformer(
            img_size=112,  # 修改为112x112
            patch_size=16,
            in_chans=3,
            num_classes=0,
            embed_dim=768,
            depth=12,
            num_heads=12,
            mlp_ratio=4.,
            qkv_bias=True,
            drop_rate=0.1,
            attn_drop_rate=0.1,
            drop_path_rate=0.1,
            norm_layer=nn.LayerNorm,
            act_layer=nn.GELU,
        )
        
        # 加载预训练的MFN模型
        self.mfn = MixedFeatureNet.MixedFeatureNet()
        if pretrained:
            self.mfn = torch.load(os.path.join('./pretrained/', "MFN_msceleb.pth"))
        
        # 特征融合层
        self.feature_fusion = nn.Sequential(
            nn.Linear(1024, 512),
            nn.LayerNorm(512),
            nn.ReLU(inplace=True)
        )
        
        # 注意力头
        self.num_head = num_head
        for i in range(int(num_head)):
            setattr(self, "cat_head%d" % (i), CoordAttHead())
        
        # 分类层
        self.Linear = Linear_block(512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0))
        self.flatten = Flatten()
        self.fc = nn.Linear(512, num_class)
        self.bn = nn.BatchNorm1d(num_class)
        
    def forward(self, x):
        # 调整输入通道
        # x = self.input_adjust(x)
        
        # ViT特征提取
        vit_features = self.vit(x)
        # vit_features = vit_features.reshape(-1, 768, 7, 7)  # 重塑为特征图
        
        # MFN特征提取
        mfn_features = self.mfn(x)
        
        # 特征融合
        combined_features = torch.cat([mfn_features, vit_features], dim=1)
        fused_features = self.feature_fusion(combined_features)
        
       
        out = self.fc(fused_features)
        
        return out

class Linear_block(Module):
    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
        super(Linear_block, self).__init__()
        self.conv = nn.Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False)
        self.bn = nn.BatchNorm2d(out_c)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

class Flatten(Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class CoordAttHead(nn.Module):
    def __init__(self):
        super().__init__()
        self.CoordAtt = CoordAtt(512, 512)
    
    def forward(self, x):
        ca = self.CoordAtt(x)
        return ca

class CoordAtt(nn.Module):
    def __init__(self, inp, oup, groups=32):
        super(CoordAtt, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))
        
        mip = max(8, inp // groups)
        
        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.conv2 = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv3 = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        identity = x
        n, c, h, w = x.size()
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)
        
        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.relu(y)
        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)
        
        x_h = self.conv2(x_h).sigmoid()
        x_w = self.conv3(x_w).sigmoid()
        x_h = x_h.expand(-1, -1, h, w)
        x_w = x_w.expand(-1, -1, h, w)
        
        y = identity * x_w * x_h
        return y